import argparse
import numpy as np
import matplotlib.pyplot as plt
from enum import Enum, auto
from typing import Tuple, List, Optional
from PIL import Image

# 设置绘图字体为黑体
plt.rcParams["font.family"] = "SimHei"


class NoiseType(Enum):
    """噪声类型枚举"""

    GAUSSIAN = auto()
    SALT_PEPPER = auto()


class EdgeDetector(Enum):
    """边缘检测算子枚举"""

    SOBEL = auto()
    KIRSCH = auto()
    LAPLACIAN = auto()
    CANNY = auto()


def read_image(path: str) -> np.ndarray:
    """手动读取并转换灰度图"""
    img = Image.open(path).convert("L")  # 使用PIL转换灰度
    return np.array(img, dtype=np.uint8)


def add_noise(img: np.ndarray, noise_type: NoiseType) -> np.ndarray:
    """手动添加噪声"""
    noisy = img.copy()
    h, w = img.shape

    if noise_type == NoiseType.GAUSSIAN:
        # 高斯噪声
        noise = np.random.normal(0, 25, (h, w))
        noisy = np.clip(noisy.astype(np.float32) + noise, 0, 255)
    elif noise_type == NoiseType.SALT_PEPPER:
        # 椒盐噪声
        prob = 0.02
        salt = np.random.rand(h, w) < prob / 2
        pepper = np.random.rand(h, w) < prob / 2
        noisy[salt] = 255
        noisy[pepper] = 0

    return noisy.astype(np.uint8)


def gaussian_kernel(size: int = 5, sigma: float = 1.0) -> np.ndarray:
    """手动生成高斯核"""
    ax = np.linspace(-(size // 2), size // 2, size)
    x, y = np.meshgrid(ax, ax)
    kernel = np.exp(-(x**2 + y**2) / (2 * sigma**2))
    return kernel / kernel.sum()


def convolve(image: np.ndarray, kernel: np.ndarray) -> np.ndarray:
    """手动实现卷积操作"""
    k_h, k_w = kernel.shape
    pad_h, pad_w = k_h // 2, k_w // 2

    # 镜像填充边界
    padded = np.pad(image, ((pad_h, pad_h), (pad_w, pad_w)), mode="reflect")

    # 初始化输出
    output = np.zeros_like(image, dtype=np.float32)

    # 滑动窗口卷积
    for i in range(image.shape[0]):
        for j in range(image.shape[1]):
            region = padded[i : i + k_h, j : j + k_w]
            output[i, j] = np.sum(region * kernel)

    return output


def median_filter(img: np.ndarray, size: int = 3) -> np.ndarray:
    """手动实现中值滤波"""
    pad = size // 2
    padded = np.pad(img, pad, mode="reflect")
    output = np.zeros_like(img)

    for i in range(img.shape[0]):
        for j in range(img.shape[1]):
            region = padded[i : i + size, j : j + size]
            output[i, j] = np.median(region)

    return output.astype(np.uint8)


def denoise(img: np.ndarray, noise_type: Optional[NoiseType]) -> np.ndarray:
    """手动实现去噪"""
    if noise_type is None:
        return img

    if noise_type == NoiseType.GAUSSIAN:
        kernel = gaussian_kernel(5, 1.0)
        return convolve(img.astype(float), kernel).clip(0, 255).astype(np.uint8)
    elif noise_type == NoiseType.SALT_PEPPER:
        return median_filter(img, 3)
    else:
        raise ValueError("不支持的噪声类型")


# 边缘检测算子实现 ============================================
def sobel_edges(img: np.ndarray) -> np.ndarray:
    """手动实现Sobel算子"""
    # 定义卷积核
    kx = np.array([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]])
    ky = np.array([[1, 2, 1], [0, 0, 0], [-1, -2, -1]])

    # 计算梯度
    gx = convolve(img.astype(float), kx)
    gy = convolve(img.astype(float), ky)

    # 计算幅值
    mag = np.sqrt(gx**2 + gy**2)
    return (mag * 255 / mag.max()).astype(np.uint8)


def kirsch_edges(img: np.ndarray) -> np.ndarray:
    """手动实现Kirsch算子"""
    kernels = [
        np.array([[-3, -3, 5], [-3, 0, 5], [-3, -3, 5]]),
        np.array([[-3, 5, 5], [-3, 0, 5], [-3, -3, -3]]),
        np.array([[5, 5, 5], [-3, 0, -3], [-3, -3, -3]]),
        np.array([[5, 5, -3], [5, 0, -3], [-3, -3, -3]]),
        np.array([[5, -3, -3], [5, 0, -3], [5, -3, -3]]),
        np.array([[-3, -3, -3], [5, 0, -3], [5, 5, -3]]),
        np.array([[-3, -3, -3], [-3, 0, -3], [5, 5, 5]]),
        np.array([[-3, -3, -3], [-3, 0, 5], [-3, 5, 5]]),
    ]

    edge_mag = np.zeros_like(img, dtype=np.float32)
    for k in kernels:
        filtered = convolve(img.astype(float), k)
        np.maximum(edge_mag, filtered, out=edge_mag)

    return (edge_mag * 255 / edge_mag.max()).astype(np.uint8)


def laplacian_edges(img: np.ndarray) -> np.ndarray:
    """手动实现Laplacian算子"""
    kernel = np.array([[0, 1, 0], [1, -4, 1], [0, 1, 0]])
    edges = convolve(img.astype(float), kernel)
    return np.abs(edges).clip(0, 255).astype(np.uint8)


def canny_edges(img: np.ndarray) -> np.ndarray:
    """手动实现Canny边缘检测"""
    # 步骤1：高斯滤波
    kernel = gaussian_kernel(5, 1.0)
    blurred = convolve(img.astype(float), kernel)

    # 步骤2：计算梯度
    kx = np.array([[-1, 0, 1]])
    ky = np.array([[-1], [0], [1]])
    gx = convolve(blurred, kx)
    gy = convolve(blurred, ky)

    # 步骤3：计算幅值和方向
    mag = np.hypot(gx, gy)
    angle = np.arctan2(gy, gx) * 180 / np.pi
    angle = np.round(angle / 45) * 45  # 量化到4个方向

    # 步骤4：非极大值抑制
    suppressed = np.zeros_like(mag)
    for i in range(1, mag.shape[0] - 1):
        for j in range(1, mag.shape[1] - 1):
            dir = angle[i, j] % 180
            # 根据方向比较相邻像素
            if dir == 0:  # 水平
                neighbors = [mag[i, j - 1], mag[i, j + 1]]
            elif dir == 45:  # 正对角线
                neighbors = [mag[i - 1, j + 1], mag[i + 1, j - 1]]
            elif dir == 90:  # 垂直
                neighbors = [mag[i - 1, j], mag[i + 1, j]]
            else:  # 反对角线
                neighbors = [mag[i - 1, j - 1], mag[i + 1, j + 1]]

            if mag[i, j] >= max(neighbors):
                suppressed[i, j] = mag[i, j]

    # 步骤5：双阈值检测
    high_thresh = suppressed.max() * 0.3
    low_thresh = high_thresh * 0.5

    strong = suppressed >= high_thresh
    weak = (suppressed >= low_thresh) & (suppressed < high_thresh)

    # 步骤6：边缘跟踪
    edges = strong.astype(np.uint8) * 255
    for i in range(1, edges.shape[0] - 1):
        for j in range(1, edges.shape[1] - 1):
            if weak[i, j]:
                if np.any(strong[i - 1 : i + 2, j - 1 : j + 2]):
                    edges[i, j] = 255

    return edges


def edge_detection(img: np.ndarray, detector: EdgeDetector) -> np.ndarray:
    """统一边缘检测接口"""
    if detector == EdgeDetector.SOBEL:
        return sobel_edges(img)
    if detector == EdgeDetector.KIRSCH:
        return kirsch_edges(img)
    if detector == EdgeDetector.LAPLACIAN:
        return laplacian_edges(img)
    if detector == EdgeDetector.CANNY:
        return canny_edges(img)
    raise ValueError("不支持的算子类型")


# 可视化与主程序 ==============================================
def plot_results(images: List[Tuple[str, np.ndarray]]):
    """结果可视化"""
    plt.figure(figsize=(20, 15))
    cols = 4
    rows = int(np.ceil(len(images) / cols))

    for i, (title, img) in enumerate(images):
        plt.subplot(rows, cols, i + 1)
        plt.imshow(img, cmap="gray")
        plt.title(title)
        plt.axis("off")

    plt.tight_layout()
    plt.show()


def main(image_path: str):
    # 读取图像
    original = read_image(image_path)

    # 准备测试数据
    test_images: List[Tuple[str, np.ndarray, Optional[NoiseType]]] = [
        ("原始图像", original, None),
        # ("高斯噪声", add_noise(original, NoiseType.GAUSSIAN), NoiseType.GAUSSIAN),
        # ("椒盐噪声", add_noise(original, NoiseType.SALT_PEPPER), NoiseType.SALT_PEPPER),
    ]

    results = []

    # 处理每个测试案例
    for title, img, noise_type in test_images:
        # 去噪处理
        denoised = denoise(img, noise_type)

        # 执行所有边缘检测
        for detector in EdgeDetector:
            # 原始/噪声图像
            edges_raw = edge_detection(img, detector)
            results.append((f"{detector.name} - {title}", edges_raw))

            # 去噪后图像
            edges_clean = edge_detection(denoised, detector)
            results.append((f"{detector.name} - {title} (去噪后)", edges_clean))

    plot_results(results)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="完全手动实现的边缘检测系统")
    parser.add_argument(
        "--image_path",
        type=str,
        default="example/作业1/lena_ 原图.jpg",
        help="输入图像路径",
    )
    args = parser.parse_args()
    if args.image_path.startswith('"') and args.image_path.endswith('"'):
        args.image_path = args.image_path[1:-1]
    print(f"参数：{args}")

    main(args.image_path)
